# -*- coding: utf-8 -*-
# backend/models/bytedance_model.py
"""
豆包模型实现
"""

from typing import Dict, Any, Optional, AsyncGenerator
from models.base_model import BaseAIModel
import aiohttp
import json

class ByteDanceModel(BaseAIModel):
    """豆包模型实现类"""
    
    def __init__(self, model_name: str, api_key: str):
        super().__init__(model_name, api_key)
        self.base_url = "https://open.bytedance.com/api/v2/chat/completions"
    
    async def initialize(self):
        """初始化模型"""
        self.session = aiohttp.ClientSession()
    
    def get_provider_name(self) -> str:
        """获取提供商名称"""
        return "字节跳动"
    
    async def generate_response(self, messages: list, **kwargs) -> Dict[str, Any]:
        """生成响应"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        params = {
            "model": "doubao-pro",
            "messages": messages,
            **kwargs
        }
        
        async with self.session.post(self.base_url, json=params, headers=headers) as response:
            response.raise_for_status()
            return await response.json()
    
    async def stream_response(self, messages: list, **kwargs) -> AsyncGenerator[str, None]:
        """流式生成响应"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        params = {
            "model": "doubao-pro",
            "messages": messages,
            "stream": True,
            **kwargs
        }
        
        async with self.session.post(self.base_url, json=params, headers=headers) as response:
            response.raise_for_status()
            async for chunk in response.content:
                if chunk:
                    yield chunk.decode('utf-8')
    
    async def close(self):
        """关闭模型"""
        await self.session.close()